import json
from zhipuai import ZhipuAI
from tqdm import tqdm  # 导入进度条库


# 加载测试集
def load_data(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return data


# 保存结果到文件
def save_results(results, file_path):
    with open(file_path, 'w', encoding='utf-8') as file:
        for result in results:
            file.write(result + '\n')


# 初始化ZhipuAI客户端
client = ZhipuAI(api_key="5baa8b8e811e40fe93dfbda226a2224a.mZqrhgAgLa1xfo9G")  # 请确保API密钥正确

# 加载测试集
test_data1 = load_data('test1.json')
test_data2 = load_data('test2.json')


# 使用智谱AI进行预测
def predict(content):
    try:
        response = client.chat.completions.create(
            model="glm-4",  # 使用GLM-4模型
            messages=[
                {
                    "role": "user",
                    "content": f"""请分析以下文本内容，并返回一个四元组结果，格式为"目标 论点 目标群体 是否仇恨":

                    文本内容: {content}

                    请按照以下格式返回结果:
                    [目标] [论点] [目标群体] [是否仇恨(是/否)]"""
                }
            ],
            temperature=0.1  # 降低随机性
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"\n预测时出错: {e}")
        return "未知 未知 未知 未知"


# 在测试集上进行预测并生成结果
def generate_results(test_data, output_file):
    results = []
    # 添加进度条
    for item in tqdm(test_data, desc="处理进度", unit="样本"):
        content = item['content']
        prediction = predict(content)
        # 清理预测结果并标准化格式
        prediction = (prediction.replace('[', '').replace(']', '').
                      replace('(', '').replace(')', ''))
        result = f"{prediction} [END]"
        results.append(result)
    save_results(results, output_file)


print("开始处理测试集1...")
generate_results(test_data1, 'test1_results.txt')

print("\n开始处理测试集2...")
generate_results(test_data2, 'test2_results.txt')

print("\n预测完成，结果已保存到文件中。")